package edu.uaskl.jqf.model.matrix;

import java.util.Arrays;
import java.util.List;

import edu.uaskl.jqf.model.Complex;
import edu.uaskl.jqf.model.Qubit;
import edu.uaskl.jqf.tools.matrix.MatrixSolvers;
import edu.uaskl.jqf.tools.matrix.MatrixTransformations;

/**
 * Immutable matrix class.<br>
 * <br>
 * A m x n matrix has m rows and n columns. <br>
 * COLUMNS<br>
 * --------><br>
 * | R<br>
 * | O<br>
 * | W<br>
 * | S<br>
 * |<br>
 * |/<br>
 * 
 * @author tbach
 */
public class MatrixComplex { // TODO add multiply scalar method, think about adding a real matrix without complex calculations
    protected final Complex[][] values;
    protected final int rows;
    protected final int columns;

    /** Creates an instance with the given rows and columns and the default value {@link Complex#get00()} */
    public MatrixComplex(final int rows, final int columns) {
        this(getArrayWithRowsColumnsDefaultValue(rows, columns, Complex.get00()), rows, columns);
    }

    /** Createsan instance with the given rows, columns and default value */
    public MatrixComplex(final int rows, final int columns, final Complex defaultValue) {
        this(getArrayWithRowsColumnsDefaultValue(rows, columns, defaultValue), rows, columns);
    }

    /** Returns new array with neutral complex number for all values */
    private static Complex[][] getArrayWithRowsColumnsDefaultValue(final int rows, final int columns, final Complex defaultValue) {
        final Complex[][] result = new Complex[rows][columns];
        for (int row = 0; row < rows; ++row)
            Arrays.fill(result[row], defaultValue);
        return result;
    }

    /** Creates a matrix backed on the given values. Changes to the given array do not can change the object behavior */
    public MatrixComplex(final Complex[][] values) {
        this(copyComplex2dArray(values), values.length, (values.length > 0) ? values[0].length : 0);
    }

    /** Creates a matrix backed on the given values. Changes to the given array can not change the object behavior */
    public MatrixComplex(final double[][] values) {
        this(Complex.convertRealToComplexArray(values), values.length, (values.length > 0) ? values[0].length : 0);
    }

    /** Creates a matrix backed on the given values. Changes to the given array <b>can change</b> the object behavior and vice versa. */
    protected MatrixComplex(final Complex[][] values, final int rows, final int columns) {
        this.values = copyComplex2dArray(values); // FIXME this should not be a copy??? -tbach
        this.rows = rows;
        this.columns = columns;
        checkValid();
        checkInitialized();
    }

    /** Creates a matrix based on the given 2d array. This means changes to the given array can change the object behavior and vice versa. */
    public static MatrixComplex createMatrixBasedOnArray2d(final Complex[][] values) {
        return new MatrixComplex(values, values.length, (values.length > 0) ? values[0].length : 0);
    }

    /** Creates a matrix based on the given 2d array. This means changes to the given array can change the object behavior and vice versa.. */
    public static MatrixComplex createMatrixBasedOnArray2d(final double[][] values) {
        return new MatrixComplex(Complex.convertRealToComplexArray(values), values.length, (values.length > 0) ? values[0].length : 0);
    }

    public Complex get(final int row, final int column) {
        return values[row][column];
    }

    public int getRows() {
        return rows;
    }

    public int getColumns() {
        return columns;
    }

    public String getDimensionAsString() {
        return new StringBuilder().append(rows).append("x").append(columns).toString();
    }

    /** Returns (a copy of) the underlying values array */
    public Complex[][] getValuesCopy() {
        return copyComplex2dArray(values);
    }

    private void checkValid() {
        for (final Complex[] row : values)
            if (row.length != columns)
                throw new IllegalStateException("Matrix is not valid. Expected row length: " + columns + ", found: " + row.length);
    }

    private void checkInitialized() {
        for (final Complex[] row : values) {
            if (row == null)
                throw new IllegalStateException("Matrix is not initialized, has null entries");
            for (final Complex field : row)
                if (field == null)
                    throw new IllegalStateException("Matrix is not initialized, has null entries");
        }
    }

    /** Returns a new matrix which is the conjugate of (this), this means all entries a are replaced by a.conj() */
    public MatrixComplex getConjugate() {
        final MatrixComplex newMatrix = new MatrixComplex(rows, columns);
        for (int row = 0; row < rows; ++row)
            for (int column = 0; column < columns; ++column)
                newMatrix.values[row][column] = values[row][column].getConjugate();
        return newMatrix;
    }

    /** Returns a new matrix which is the transpose of (this), this means a_ij goes to a_ij for all indices i and j */
    public MatrixComplex getTranspose() {
        final MatrixComplex newMatrix = new MatrixComplex(columns, rows);
        for (int row = 0; row < rows; ++row)
            for (int column = 0; column < columns; ++column)
                newMatrix.values[column][row] = values[row][column]; // not nice for cache access, TODO implement view? (tbach)
        return newMatrix;
    }

    /** Returns a new matrix which is the conjugate transpose of (this) */
    public MatrixComplex getAdjoint() {
        final MatrixComplex newMatrix = new MatrixComplex(columns, rows);
        for (int row = 0; row < rows; ++row)
            for (int column = 0; column < columns; ++column)
                newMatrix.values[column][row] = values[row][column].getConjugate(); // not nice for cache access, TODO implement view? (tbach)
        return newMatrix;
    }

    public MatrixComplex multiply(final MatrixComplex other) {
        if (columns != other.rows)
            throw new IllegalArgumentException("Matrices do not match: columns=" + columns + " != other.rows=" + other.rows);

        final Complex[][] newValues = getArrayWithRowsColumnsDefaultValue(rows, other.columns, Complex.get00());
        for (int row = 0; row < rows; ++row)
            for (int column = 0; column < columns; column++)
                // we could create a sum variable, but it does not change the overall operations (tbach)
                for (int otherColumn = 0; otherColumn < other.columns; ++otherColumn)
                    newValues[row][otherColumn] = newValues[row][otherColumn].add(values[row][column].multiply(other.values[column][otherColumn]));
        return createMatrixBasedOnArray2d(newValues);
    }

    public Qubit multiply(final Qubit other) {
        final MatrixComplex newMatrix = multiply((MatrixComplex) other);
        final Complex newAlpha = newMatrix.get(0, 0);
        final Complex newBeta = newMatrix.get(1, 0);
        return new Qubit(newAlpha, newBeta);
    }

    public MatrixComplex multiply(final Complex value) {
        final Complex[][] newValues = new Complex[rows][columns];
        for (int row = 0; row < rows; ++row)
            for (int column = 0; column < columns; column++)
                newValues[row][column] = values[row][column].multiply(value);
        return createMatrixBasedOnArray2d(newValues);
    }

    public MatrixComplex multiply(final double value) {
        return multiply(new Complex(value));
    }

    /** Returns if adj(this) * this = I */
    public boolean isUnitary() {
        return this.getAdjoint().multiply(this).isIdentity();
    }

    /**
     * Returns the tensor/kronecker product of (this) and other<br>
     * See: <a href=http://de.wikipedia.org/wiki/Kronecker-Produkt>http://de.wikipedia.org/wiki/Kronecker-Produkt</a>
     * Input sizes: A of size m x n, B of size p x q <br>
     * Output sizes: A tensor B of size mp x nq <br>
     */
    public MatrixComplex tensor(final MatrixComplex other) {
        // input sizes: A of size m x n, B of size p x q
        // output sizes: A x B of size mp x nq
        final Complex[][] newValues = new Complex[rows * other.rows][columns * other.columns];
        for (int row = 0; row < rows; ++row)
            for (int column = 0; column < columns; ++column) {
                final int rowField = row * other.rows;
                final int columnField = column * other.columns;
                final Complex factor = values[row][column];
                for (int otherRow = 0; otherRow < other.rows; ++otherRow)
                    for (int otherColumn = 0; otherColumn < other.columns; ++otherColumn)
                        newValues[rowField + otherRow][columnField + otherColumn] = factor.multiply(other.values[otherRow][otherColumn]);
            }
        return createMatrixBasedOnArray2d(newValues);
    }

    /**
     * Returns the partial trace of the given size (size of the base room) for the first base room<br>
     * For example if we have a matrix A with a 2x2 size and a matrix B with a 4x4 size,
     * then the tensor ({@link #tensor(MatrixComplex)}) A tensor B = C is of size 8x8.<br>
     * Now, if we calculate the first partial trace Trp(C) with size 2, we will get back a 2x2 matrix A'.<br>
     * A' = x * A, with a scalar x. x = 1 if B is a valid density matrix (because we trace out A*tr(B)).<br>
     * This means we get back A, if B is a valid density matrix of a state, otherwise is is not guaranteed that the trace is 1.<br>
     * <br>
     * Current version only works for a quadratic matrix and the argument size must divide the size of the matrix.<br>
     * If the second argument (in the example: B) should be traced, we have to use {@link #getPartialTraceSecond(int)}.
     */
    public MatrixComplex getPartialTraceFirst(final int size) { // Returns A
        if (!isQuadratic())
            throw new IllegalStateException("matrix is not quadratic. Algorithm implemented only for quadratic matrix. current matrix: "
                    + getDimensionAsString());
        if ((size <= 0) || ((columns % size) != 0))
            throw new IllegalArgumentException("argument size must divide the size of the matrix. matrix: " + getDimensionAsString() + ", argument size: "
                    + size);
        /*
         * Idea: Go trough all diagonal fields, sum them up to a new matrix.
         * Example: A=2x2 (size=2), B=4x4, A tensor B = C=8x8 matrix:
         * ########
         * ########
         * ########
         * ########
         * ########
         * ########
         * ########
         * ########
         * 
         * Get the diagonal fields:
         * a###a###
         * #b###b##
         * ##c###c#
         * ###d###d
         * a###a###
         * #b###b##
         * ##c###c#
         * ###d###d
         * 
         * a,b,c,d is 2x2
         * Matrix a+b+c+d=A, because of a+b+c+d = (b11+b22+b33+b44)*A and (b11+b22+b33+b44)=1 because B is normalized
         * 
         */
        final Complex[][] newValues = getArrayWithRowsColumnsDefaultValue(size, size, Complex.get00());
        final int fieldSize = columns / size;
        for (int i = 0; i < fieldSize; i++)
            for (int rowField = 0; rowField < size; ++rowField) {
                final int row = (rowField * fieldSize) + i;
                for (int columnField = 0; columnField < size; ++columnField) {
                    final int column = (columnField * fieldSize) + i;
                    newValues[rowField][columnField] = newValues[rowField][columnField].add(values[row][column]);
                }

            }
        return createMatrixBasedOnArray2d(newValues);
    }

    /**
     * Returns the partial trace of the given size (size of the base room) for the second base room<br>
     * For example if we have a matrix A with a 2x2 size and a matrix B with a 4x4 size,
     * then the tensor ({@link #tensor(MatrixComplex)}) A tensor B = C is of size 8x8.<br>
     * Now, if we calculate the second partial trace Trp(C) with size 4, we will get back a 4x4 matrix B'.<br>
     * B' = x * B, with a scalar x. x = 1 if A is a valid density matrix (because we trace out B*tr(A)).<br>
     * This means we get back B, if A is a valid density matrix of a state, otherwise is is not guaranteed that the trace is 1.<br>
     * <br>
     * Current version only works for a quadratic matrix and the argument size must divide the size of the matrix.<br>
     * If the first argument (in the example: A) should be traced, we have to use {@link #getPartialTraceFirst(int)}.
     */
    public MatrixComplex getPartialTraceSecond(final int size) { // Returns B
        if (!isQuadratic())
            throw new IllegalStateException("matrix is not quadratic. Algorithm implemented only for quadratic matrix. current matrix: "
                    + getDimensionAsString());
        if ((size <= 0) || ((columns % size) != 0))
            throw new IllegalArgumentException("argument size must divide the size of the matrix. matrix: " + getDimensionAsString() + ", argument size: "
                    + size);
        /*
         * Idea: Go trough all diagonal fields, sum them up to a new matrix.
         * Example: A=2x2, B=4x4 (size=4), A tensor B = C=8x8 matrix:
         * ########
         * ########
         * ########
         * ########
         * ########
         * ########
         * ########
         * ########
         * 
         * Get the diagonal fields:
         * mmmm####
         * mmmm####
         * mmmm####
         * mmmm####
         * ####nnnn
         * ####nnnn
         * ####nnnn
         * ####nnnn
         * 
         * Matrix m+n=B, because of m+n = (a11 + a22)*B and a11+a22=1 because A is normalized
         * 
         */
        final Complex[][] newValues = getArrayWithRowsColumnsDefaultValue(size, size, Complex.get00());
        for (int currentField = 0; currentField < rows; currentField += size)
            for (int row = 0; row < size; ++row)
                for (int column = 0; column < size; ++column)
                    newValues[row][column] = newValues[row][column].add(values[currentField + row][currentField + column]);

        return createMatrixBasedOnArray2d(newValues);
    }

    /** Returns the trace, this is the sum of all diagonal elements. */
    public Complex getTrace() {
        final int numberOfDiagonalElements = Math.min(rows, columns);
        Complex result = Complex.get00();
        for (int i = 0; i < numberOfDiagonalElements; i++)
            result = result.add(values[i][i]);
        return result;
    }

    /** Checks if all diagonal entries are the neutral element of multiplication and rest is neutral element of addition */
    public boolean isIdentity() {
        if (!isQuadratic())
            return false;
        for (int row = 0; row < rows; ++row)
            for (int column = 0; column < columns; ++column) {
                final Complex currentValue = values[row][column];
                if (row == column) {
                    if (!currentValue.isNeutralToMultiplication())
                        return false;
                } else if (!currentValue.isNeutralToAddition())
                    return false;
            }
        return true;
    }

    public boolean isQuadratic() {
        return rows == columns;
    }

    public boolean isSize(final int columnsOther, final int rowsOther) {
        return (this.columns == columnsOther) && (this.rows == rowsOther);
    }

    public List<Complex> getEigenValues() {
        final MatrixTransformationHessenberg matrixTransformationHessenberg = MatrixTransformations.HESSENBERG.transform(this);
        final MatrixSolverEigen matrixSolverEigen = MatrixSolvers.EIGEN.solve(matrixTransformationHessenberg);
        return matrixSolverEigen.getEigenValues();
    }

    /**
     * Calculates the entropy of the this system:<br>
     * Sum over all eigenvalues y: y * log(y)<br>
     * We use {@link Complex#getLog2()} in this implementation.
     * This conforms to Nielsen, Chuang - Quantum Computation and Quantum Information, 2010, P510 Def. 11.40
     * and McMahon - Quantum Computing Explained 2008, P298, Def 13.39
     */
    public Complex getEntropy() {
        return getEntropyToLogBase(2);
    }

    /**
     * Calculates the entropy of the this system:<br>
     * Sum over all eigenvalues y: y * log_base(y)<br>
     * We use {@link Complex#getLogToBase(double)} in this implementation.
     */
    public Complex getEntropyToLogBase(final double base) {
        final List<Complex> eigenValues = getEigenValues();
        Complex entropy = Complex.get00();
        // entropy = sum over all eigenvalues y: -y*log(y) -tbach
        for (final Complex complex : eigenValues) {
            if (complex.equals(Complex.get00())) // or we have log(0) = NaN. This is ok to skip, because limit x->0 for x*log(x) = 0 -tbach
                continue;
            entropy = entropy.add(complex.multiply(complex.getLogToBase(base)));
        }
        return entropy.multiply(new Complex(-1));
    }

    private static Complex[][] copyComplex2dArray(final Complex[][] array) {
        final Complex[][] result = new Complex[array.length][];
        for (int i = 0; i < array.length; i++)
            result[i] = Arrays.copyOf(array[i], array[i].length);
        return result;
    }

    @Override
    public int hashCode() {
        final int prime = 31;
        int result = 1;
        result = (prime * result) + columns;
        result = (prime * result) + rows;
        result = (prime * result) + Arrays.hashCode(values);
        return result;
    }

    @Override
    public boolean equals(final Object obj) {
        if (this == obj)
            return true;
        if (obj == null)
            return false;
        if (getClass() != obj.getClass())
            return false;
        final MatrixComplex other = (MatrixComplex) obj;
        if (columns != other.columns)
            return false;
        if (rows != other.rows)
            return false;
        if (!Arrays.deepEquals(values, other.values))
            return false;
        return true;
    }

    @Override
    public String toString() {
        final StringBuilder stringBuilder = new StringBuilder();
        int longestComplexString = 0;
        for (final Complex[] row : values)
            for (final Complex rowEntry : row)
                longestComplexString = Math.max(longestComplexString, rowEntry.toString().length());
        final String formatString = " %" + longestComplexString + "s |";
        for (final Complex[] row : values) {
            stringBuilder.append("|");
            for (final Complex rowEntry : row)
                stringBuilder.append(String.format(formatString, rowEntry));
            stringBuilder.append("\n");
        }
        return stringBuilder.toString();
    }
}
